//
// Created by DELL on 2022/4/20.
//

#include "Pipeline.h"
#include "Shader.h"

#include <utility>
#include "Device.h"

BEGIN_NS_VULKAN

Pipeline::Pipeline(DeviceRef device)
    : device(std::move(device))
{

}

Pipeline::~Pipeline()
{
    vkDestroyPipeline(device->getDevice(), graphicsPipeline, nullptr);
    vkDestroyPipelineLayout(device->getDevice(), pipelineLayout, nullptr);
}

PipelineRef Pipeline::create(DeviceRef& device, const ShaderRef& shader, const PipelineConfigInfo& configInfo)
{
    auto pipeline = new Pipeline(device);
    pipeline->init(shader, configInfo);

    return Vulkan::PipelineRef(pipeline);
}

void Pipeline::init(const ShaderRef& shader, const PipelineConfigInfo& configInfo)
{
    std::vector<VkShaderModule> shaderModules(shader->getStageCount());
    std::vector<VkPipelineShaderStageCreateInfo> shaderStates(shader->getStageCount());
    for (int i = 0; i < shader->getStageCount(); ++i)
    {
        createShaderModule(shader->getSpirv(i), &shaderModules[i]);
        shaderStates[i].sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
        shaderStates[i].stage = shader->getStage(i);
        shaderStates[i].module = shaderModules[i];
        shaderStates[i].pName = "main";
        shaderStates[i].flags = 0;
        shaderStates[i].pNext = nullptr;
        shaderStates[i].pSpecializationInfo = nullptr;
    }

    VkPipelineVertexInputStateCreateInfo vertexInputInfo{};
    vertexInputInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
    vertexInputInfo.vertexBindingDescriptionCount = static_cast<uint32_t>(configInfo.vertexBindings.size());
    vertexInputInfo.pVertexBindingDescriptions = configInfo.vertexBindings.data();
    vertexInputInfo.vertexAttributeDescriptionCount = static_cast<uint32_t>(configInfo.vertexAttributes.size());
    vertexInputInfo.pVertexAttributeDescriptions = configInfo.vertexAttributes.data();

    VkPipelineViewportStateCreateInfo viewportInfo{};
    viewportInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
    viewportInfo.viewportCount = 1;
    viewportInfo.pViewports = nullptr;
    viewportInfo.scissorCount = 1;
    viewportInfo.pScissors = nullptr;

    VkPipelineColorBlendStateCreateInfo colorBlendInfo{};
    colorBlendInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
    colorBlendInfo.logicOpEnable = VK_FALSE;
    colorBlendInfo.logicOp = VK_LOGIC_OP_COPY;
    colorBlendInfo.attachmentCount = 1;
    colorBlendInfo.pAttachments = &configInfo.colorBlendAttachment;
    colorBlendInfo.blendConstants[0] = 0.0f;
    colorBlendInfo.blendConstants[1] = 0.0f;
    colorBlendInfo.blendConstants[2] = 0.0f;
    colorBlendInfo.blendConstants[3] = 0.0f;

    VkPipelineDynamicStateCreateInfo dynamicStateInfo{};
    dynamicStateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
    dynamicStateInfo.pDynamicStates = configInfo.dynamicStateEnables.data();
    dynamicStateInfo.dynamicStateCount = static_cast<uint32_t>(configInfo.dynamicStateEnables.size());
    dynamicStateInfo.flags = 0;

    VkPipelineLayoutCreateInfo pipelineLayoutInfo{};
    pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
    pipelineLayoutInfo.setLayoutCount = static_cast<uint32_t>(configInfo.setLayouts.size());
    pipelineLayoutInfo.pSetLayouts = configInfo.setLayouts.data();
    pipelineLayoutInfo.pushConstantRangeCount = 0;
    pipelineLayoutInfo.pPushConstantRanges = nullptr;
    if (vkCreatePipelineLayout(device->getDevice(), &pipelineLayoutInfo, nullptr, &pipelineLayout))
    {
        throw std::runtime_error("failed to create pipeline layout.");
    }

    VkGraphicsPipelineCreateInfo pipelineInfo{};
    pipelineInfo.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
    pipelineInfo.stageCount = static_cast<uint32_t>(shaderStates.size());
    pipelineInfo.pStages = shaderStates.data();
    pipelineInfo.pVertexInputState = &vertexInputInfo;
    pipelineInfo.pInputAssemblyState = &configInfo.inputAssemblyInfo;
    pipelineInfo.pViewportState = &viewportInfo;
    pipelineInfo.pRasterizationState = &configInfo.rasterizationInfo;
    pipelineInfo.pMultisampleState = &configInfo.multisampleInfo;
    pipelineInfo.pColorBlendState = &colorBlendInfo;
    pipelineInfo.pDynamicState = &dynamicStateInfo;
    pipelineInfo.pDepthStencilState = &configInfo.depthStencilInfo;

    pipelineInfo.layout = pipelineLayout;
    pipelineInfo.renderPass = configInfo.renderPass;
    pipelineInfo.subpass = configInfo.subpass;

    pipelineInfo.basePipelineHandle = VK_NULL_HANDLE;
    pipelineInfo.basePipelineIndex = -1;

    if (vkCreateGraphicsPipelines(
            device->getDevice(),
            VK_NULL_HANDLE,
            1,
            &pipelineInfo,
            nullptr,
            &graphicsPipeline) != VK_SUCCESS)
    {
        throw std::runtime_error("failed to create graphics pipeline!");
    }
    for (auto sm : shaderModules)
    {
        vkDestroyShaderModule(device->getDevice(), sm, nullptr);
    }
}

void Pipeline::defaultPipelineConfigInfo(PipelineConfigInfo& info)
{

    info.inputAssemblyInfo.sType                  = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
    info.inputAssemblyInfo.topology               = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;
    info.inputAssemblyInfo.primitiveRestartEnable = VK_FALSE;

    info.rasterizationInfo.sType                   = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
    info.rasterizationInfo.depthClampEnable        = VK_FALSE;
    info.rasterizationInfo.rasterizerDiscardEnable = VK_FALSE;
    info.rasterizationInfo.polygonMode             = VK_POLYGON_MODE_FILL;
    info.rasterizationInfo.lineWidth               = 1.0f;
    info.rasterizationInfo.cullMode                = VK_CULL_MODE_NONE;
    info.rasterizationInfo.frontFace               = VK_FRONT_FACE_CLOCKWISE;
    info.rasterizationInfo.depthBiasEnable         = VK_FALSE;
    info.rasterizationInfo.depthBiasConstantFactor = 0.0f;
    info.rasterizationInfo.depthBiasClamp          = 0.0f;
    info.rasterizationInfo.depthBiasSlopeFactor    = 0.0f;

    info.multisampleInfo.sType                 = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
    info.multisampleInfo.sampleShadingEnable   = VK_FALSE;
    info.multisampleInfo.rasterizationSamples  = VK_SAMPLE_COUNT_1_BIT;
    info.multisampleInfo.minSampleShading      = 1.0f;
    info.multisampleInfo.pSampleMask           = nullptr;
    info.multisampleInfo.alphaToCoverageEnable = VK_FALSE;
    info.multisampleInfo.alphaToOneEnable      = VK_FALSE;
    info.multisampleInfo.flags = 0;

    info.colorBlendAttachment.colorWriteMask =
        VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT | VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;
    info.colorBlendAttachment.blendEnable         = VK_FALSE;
    info.colorBlendAttachment.srcColorBlendFactor = VK_BLEND_FACTOR_ONE;
    info.colorBlendAttachment.dstColorBlendFactor = VK_BLEND_FACTOR_ZERO;
    info.colorBlendAttachment.colorBlendOp        = VK_BLEND_OP_ADD;
    info.colorBlendAttachment.srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE;
    info.colorBlendAttachment.dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO;
    info.colorBlendAttachment.alphaBlendOp        = VK_BLEND_OP_ADD;

    info.depthStencilInfo.sType                 = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO;
    info.depthStencilInfo.depthTestEnable       = VK_FALSE;
    info.depthStencilInfo.depthWriteEnable      = VK_TRUE;
    info.depthStencilInfo.depthCompareOp        = VK_COMPARE_OP_LESS;
    info.depthStencilInfo.depthBoundsTestEnable = VK_FALSE;
    info.depthStencilInfo.minDepthBounds        = 0.0f;
    info.depthStencilInfo.maxDepthBounds        = 1.0f;
    info.depthStencilInfo.stencilTestEnable     = VK_FALSE;
    info.depthStencilInfo.front                 = {};
    info.depthStencilInfo.back                  = {};
    info.depthStencilInfo.flags = 0;

    info.dynamicStateEnables = {VK_DYNAMIC_STATE_VIEWPORT, VK_DYNAMIC_STATE_SCISSOR};
}

void Pipeline::createShaderModule(const std::vector<uint32_t>& code, VkShaderModule* shaderModule)
{
    VkShaderModuleCreateInfo createInfo{};
    createInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
    createInfo.codeSize = code.size() * sizeof(uint32_t);
    createInfo.pCode = reinterpret_cast<const uint32_t*>(code.data());

    if (vkCreateShaderModule(device->getDevice(), &createInfo, nullptr, shaderModule) != VK_SUCCESS)
    {
        throw std::runtime_error("failed to create shader module.");
    }
}


END_NS_VULKAN
